# 编写此程序的目的在于在确定网络拓扑结构以及网络设备硬件各项性能之后
# 确定主机端到端进行通信时的最大发送速率，TCP报文的发送窗口大小，以及udp报文的大小
# 使得能够充分运用网络带宽，提高网络传输效率
import subprocess
import pandas as pd
import argparse

def run_iperf(server_ip="localhost",port=5001, duration=10,**kwargs):
    # 构建iperf命令
    command = ['iperf', '-c', server_ip, '-p', str(port), '-t', str(duration),'-f','m','-u']
    if "bandwidth" in kwargs:
        command.extend(['-b', str(kwargs["bandwidth"])])
    if "window_size" in kwargs:
        command.extend(['-w', str(kwargs["window_size"])])
    if "datagram_length" in kwargs:
        command.extend(['-l', str(kwargs["datagram_length"])])
    if "time" in kwargs:
        command.extend(['-t', str(kwargs["time"])])

    try:
        # 执行命令并捕获输出
        result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        # 检查命令是否成功执行
        if result.returncode == 0:
            pass
        else:
            print("iperf command failed.")
            print("Error:\n", result.stderr)
        
        # 返回标准输出和标准错误
        return result.stdout, result.stderr
    
    except Exception as e:
        print(f"An error occurred when execute iperf: {e}")
        return None, None

# 创建dadaframe记录测试结果
df = pd.DataFrame(columns=["datagram_length", "window_size", "bandwidth", "time", "jitter","throughput", "lost_rate"])

def interact_with_user():
    argparser = argparse.ArgumentParser(description="iperf test tool")
    argparser.add_argument("-s", "--server", help="Server IP address", default="localhost")
    argparser.add_argument("-p", "--port", help="Server port", default=5001)
    argparser.add_help("-s --server\t test server ip address\n -p --port\t test server port")
    args = argparser.parse_args()
    server_ip = args.server
    server_port = args.port

    f=open("to_%s_iperf_report.txt" % server_ip,"w")
    while True:
        try:
            datagram_length, window_size, bandwidth, time = input("input 0 to bypass paticular argument, Datagram length, Window size, Bandwidth, Time: ").split()
            kwargs={}
            if datagram_length != "0":
                kwargs["datagram_length"] = datagram_length
            if window_size != "0":
                kwargs["window_size"] = window_size
            if bandwidth != "0":
                kwargs["bandwidth"] = bandwidth
            if time != "0":
                kwargs["time"] = time
            # 运行iperf命令
            stdout, stderr = run_iperf(server_ip=server_ip, port=server_port, **kwargs)
            if stdout:
                # 从输出中提取吞吐量和丢包率
                output = stdout.decode("utf-8")
                result=output.split("\n")[-2]
                f.write(result)
                print(result,flush=True)
                throughput = float(result.split()[-7])
                lost_rate = float(result.split()[-1][1:-2]) / 100
                jitter = float(result.split()[-5])

                # 将测试结果添加到dataframe中
                df.loc[len(df)] = [datagram_length, window_size, bandwidth, time, jitter,throughput, lost_rate]
        except Exception as e:
            print(f"An error occurred: {e}")
            break
        except KeyboardInterrupt as e:
            # 测试结束
            f.close()
            df.to_csv("to_%s_iperf_report.csv" % server_ip,columns=["datagram_length", "window_size", "bandwidth", "time", "jitter","throughput", "lost_rate"],mode='w')

interact_with_user()